#!/usr/bin/env python3
"""Quick test to verify vLLM and Ollama connectivity."""

import asyncio
import sys
sys.path.insert(0, '..')

from src.config import VLLMConfig, OllamaConfig
from src.clients import VLLMClient, OllamaClient
import aiohttp


async def test_vllm():
    """Test vLLM connection and simple generation."""
    print("\n🔍 Testing vLLM...")
    config = VLLMConfig()
    client = VLLMClient(config)
    
    try:
        async with aiohttp.ClientSession() as session:
            result = await client.generate(
                session,
                "你好，请简单介绍一下水文学。",
                max_tokens=50
            )
            
            if result.success:
                print(f"✅ vLLM connection successful!")
                print(f"   Response: {result.response[:100]}...")
                print(f"   Latency: {result.latency:.2f}s")
                print(f"   Tokens: {result.tokens}")
            else:
                print(f"❌ vLLM test failed: {result.error}")
    except Exception as e:
        print(f"❌ vLLM test exception: {e}")


async def test_ollama():
    """Test Ollama connection and simple generation."""
    print("\n🔍 Testing Ollama...")
    config = OllamaConfig()
    client = OllamaClient(config)
    
    try:
        async with aiohttp.ClientSession() as session:
            result = await client.generate(
                session,
                "你好，请简单介绍一下水文学。",
                max_tokens=50
            )
            
            if result.success:
                print(f"✅ Ollama connection successful!")
                print(f"   Response: {result.response[:100]}...")
                print(f"   Latency: {result.latency:.2f}s")
                print(f"   Tokens: {result.tokens}")
            else:
                print(f"❌ Ollama test failed: {result.error}")
    except Exception as e:
        print(f"❌ Ollama test exception: {e}")


async def main():
    print("="*60)
    print("🧪 Quick Connectivity Test")
    print("="*60)
    
    await test_vllm()
    await test_ollama()
    
    print("\n" + "="*60)
    print("✅ Test Complete!")
    print("="*60)


if __name__ == "__main__":
    asyncio.run(main())

